from torch_geometric.datasets import TUDataset,Planetoid,Amazon,Coauthor,Reddit
from torch_geometric.loader import DataLoader
import os
import wandb
import dgl
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
# 获取上一级目录

import dgl
from torch_geometric.loader import DataLoader
from torch_geometric.data import Data

import os
import dgl
import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader

def dgl_to_tg(dgl_graph):
    """将 DGL Graph 转换为 PyTorch Geometric 的 Data 格式"""
    edge_index = torch.stack(dgl_graph.edges())  # 获取边索引
    x = dgl_graph.ndata['feat']  # 获取节点特征
    # y = dgl_graph.ndata['label']  # 获取标签
    # return Data(x=x, edge_index=edge_index, y=y)
    return Data(x=x, edge_index=edge_index)

def get_loader_pretrain_data(dataset):
    """从 bin 文件加载 DGL Graph，并转换为 PyG DataLoader"""
    
    # 预定义数据集名称（bin 文件名）
    datasets = ['Cora', 'Citeseer', 'Pubmed', 'Photo', 'Computers', 'Cora_1', 'Cora_2', 'Cora_3', 'Cora_4', 'Cora_5']
    
    # 获取数据存储路径
    data_path = wandb.config.bin_data_path  # 确保 wandb.config 里有 bin_data_path

    # 存储 DataLoader
    loaders = []
    
    
    for dataset in datasets:
        bin_file = os.path.join(data_path, f"{dataset}.bin")
        
        # 加载 DGL 图
        g_list, _ = dgl.load_graphs(bin_file)  # 可能有多个图，这里取第一个
        dgl_graph = g_list[0]

        # 转换为 PyG Data 格式
        tg_graph = dgl_to_tg(dgl_graph)

        # 创建 PyG DataLoader
        loader = DataLoader([tg_graph], batch_size=1)
        
        loaders.append(loader)

    return loaders[0], loaders[2], loaders[3], loaders[4], loaders[5], loaders[6], loaders[7], loaders[8], loaders[9]




def get_loader_down_data(dataset):
    # 给我下游你的数据名称 我给你数据集
    config = wandb.config
    data_path = config.data_path
    if dataset == 'Cora':
        data_down = Planetoid(root=data_path, name='Cora')
        loader = DataLoader(data_down)
             
    elif dataset == 'Citeseer':
        data_down = Planetoid(root=data_path, name='Citeseer')
        loader = DataLoader(data_down)
        
    elif dataset == 'Pubmed':
        data_down = Planetoid(root=data_path, name='Pubmed')
        loader = DataLoader(data_down)
        
    elif dataset == 'Photo':
        data_down = Amazon(root=data_path, name='Photo')
        loader = DataLoader(data_down)
        
    elif dataset == 'Computers':
        data_down = Amazon(root=data_path, name='Computers')
        loader = DataLoader(data_down)
        
    elif dataset == 'Reddit':
        data_path1 = data_path
        data_path1 = os.path.join(data_path1, 'Reddit')
        data_down = Reddit(root=data_path1)
        loader = DataLoader(data_down)
        
    else:
        raise ValueError(f"Dataset {dataset} is not recognized.")
    
    return loader
    